import os
import uuid
import json
from http import HTTPStatus

import requests
import cv2
from typing import Tuple, Dict, Any

from dashscope.common.error import InvalidTask
from dotenv import load_dotenv
from openai import OpenAI

from utils import sr_s3
from utils.sr_utils import image_path_to_data_url, content_moderations, video_compress_image_size_path
from dashscope import VideoSynthesis


# import dashscope

# dashscope.base_http_api_url = 'https://dashscope-intl.aliyuncs.com/api/v1'
#
from utils.status_codes import StatusCodes


class wanXiangGenerator_Vid:
    """视频生成工具类，封装KlingAI视频生成API调用逻辑，支持外部调用"""

    def __init__(self):
        """初始化配置与环境变量"""
        # 加载环境变量
        self.verify_ssl = True
        load_dotenv(verbose=True, override=True)
        self.app_env = os.getenv('APP_ENV', 'prod')
        self.url = "https://api.vidu.com"
        self.API_KEY = os.getenv('WXKEY', "sk-0956b1646b5f438185de690479d03e26")
        self.NSFW_CHECK_ENABLED = os.getenv('NSFW_CHECK_ENABLED')
        # 初始化鉴权头（JWT动态生成，30分钟有效期）

        # 临时文件存储目录（可外部配置）
        self.temp_image_dir = "./upload"
        self.temp_video_dir = "./download"
        os.makedirs(self.temp_image_dir, exist_ok=True)
        os.makedirs(self.temp_video_dir, exist_ok=True)

        # 分辨率映射：{档位: {宽高比: 分辨率, "default": 默认分辨率}}
        self.RESOLUTION_MAP = {
            "480P": {
                "16:9": "832*480",
                "9:16": "480*832",
                "1:1": "624*624",
                "default": None  # 480P无全局默认，由模型决定
            },
            "720P": {
                "16:9": "1280*720",
                "9:16": "720*1280",
                "1:1": "960*960",
                "4:3": "1088*832",
                "3:4": "832*1088",
                "default": "1280*720"  # 720P默认分辨率
            },
            "1080P": {
                "16:9": "1920*1080",
                "9:16": "1080*1920",
                "1:1": "1440*1440",
                "4:3": "1632*1248",
                "3:4": "1248*1632",
                "default": "1920*1080"  # 1080P默认分辨率
            }
        }

        # 模型支持的档位映射：{模型名称: 支持的档位列表}
        self.MODEL_SUPPORTED_QUALITIES = {
            "wan2.2-t2v-plus": ["480P", "1080P"],
            "wanx2.1-t2v-turbo": ["480P", "720P"],
            "wanx2.1-t2v-plus": ["720P"]
        }

    def get_resolution(self, model: str, quality: str, aspect_ratio: str = None) -> tuple[bool, str]:
        """
        根据模型、分辨率档位、宽高比获取具体分辨率

        参数:
            model: 模型名称（如"wan2.2-t2v-plus"）
            quality: 分辨率档位（如"480P"、"720P"、"1080P"）
            aspect_ratio: 宽高比（如"16:9"，为None时返回默认分辨率）

        返回:
            (是否成功, 分辨率字符串/错误信息)
        """
        # 1. 校验模型是否支持该档位
        if model not in self.MODEL_SUPPORTED_QUALITIES:
            return False, f"不支持的模型: {model}"
        if quality not in self.MODEL_SUPPORTED_QUALITIES[model]:
            return False, f"模型{model}不支持{quality}档位"

        # 2. 校验档位是否存在
        if quality not in self.RESOLUTION_MAP:
            return False, f"无效的分辨率档位: {quality}"

        # 3. 处理分辨率获取（优先按宽高比，无则用默认）
        quality_map = self.RESOLUTION_MAP[quality]
        if aspect_ratio:
            if aspect_ratio in quality_map:
                return True, quality_map[aspect_ratio]
            else:
                return False, f"{quality}档位不支持{aspect_ratio}宽高比"
        else:
            # 无宽高比时返回默认分辨率（特殊处理480P的默认）
            default_res = quality_map["default"]
            if default_res:
                return True, default_res
            else:
                # 480P无默认，需强制指定宽高比
                return False, f"{quality}档位需指定宽高比（支持: {list(quality_map.keys())[:-1]}）"

    # --------------------------
    # 内部工具方法：鉴权与配置
    # --------------------------

    # def _encode_jwt_token(self):
    #     headers = {
    #         "alg": "HS256",
    #         "typ": "JWT"
    #     }
    #     payload = {
    #         "iss": "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BAp",
    #         "exp": int(time.time()) + 1800,  # 有效时间，此处示例代表当前时间+1800s(30min)
    #         "nbf": int(time.time()) - 5  # 开始生效的时间，此处示例代表当前时间-5秒
    #     }
    #     token = jwt.encode(payload, "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BAp", headers=headers)
    #     return token

    # --------------------------
    # 内部工具方法：文件处理
    # --------------------------
    def _delete_temp_files(self, *file_paths: str) -> None:
        """删除临时文件（忽略不存在的文件）"""
        for path in file_paths:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"清理临时文件失败 {path}: {str(e)}")

    def _download_and_compress_image(self, img_url: str, storage_service: str, bucket_name: str) -> Tuple[bool, str, int]:

        from utils import sr_s3
        """下载图片并压缩（返回是否成功、base64编码、是否压缩）"""

        if self.NSFW_CHECK_ENABLED == "true":
            image_url = sr_s3.get_url(img_url, serviceName=storage_service, bucketName=bucket_name)
            clientOpenAI = OpenAI()
            image_result = content_moderations(clientOpenAI, image_url)
            if image_result:
                return False, f"图片违规", 5
        try:
            # 生成临时保存路径
            name = str(uuid.uuid4())
            ext = os.path.splitext(img_url)[-1].lower() or '.jpg'
            save_path = os.path.join(self.temp_image_dir, f"{name}{ext}")
            # 调用S3工具下载图片（假设sr_s3为外部存储工具）
            success, err = sr_s3.download_file(
                img_url, save_path, serviceName=storage_service, bucketName=bucket_name
            )
            if not success:
                return False, f"下载失败: {err}", 2
            base64_img = image_path_to_data_url(save_path)
            return True, base64_img, 1
        except Exception as e:
            return False, f"图片处理异常: {str(e)}", 2

    def _download_video(self, video_url: str, task_id: str) -> Tuple[bool, Dict[str, Any]]:
        """下载生成的视频到本地（返回是否成功、视频信息）"""
        try:
            filename = f"video_{task_id}.mp4"
            save_path = os.path.join(self.temp_video_dir, filename)

            # 流式下载视频
            response = requests.get(video_url, stream=True, timeout=30)
            response.raise_for_status()

            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB/块
                    if chunk:
                        f.write(chunk)

            # 获取视频信息（宽高、帧率等）
            cap = cv2.VideoCapture(save_path)
            video_info = {
                "save_path": save_path,
                "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "fps": cap.get(cv2.CAP_PROP_FPS),
                "duration": cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0
            }
            cap.release()
            return True, video_info
        except Exception as e:
            return False, {"error": f"视频下载失败: {str(e)}"}

    def _upload_video_to_s3(self, video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        """上传视频到S3（返回是否成功、S3路径）"""
        if not video_path or not os.path.exists(video_path):
            return False, "视频文件不存在"
        try:
            from utils import sr_s3
            s3_key = sr_s3.upload_login_or_visitor(
                video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
            )
            return True, s3_key
        except Exception as e:
            return False, f"上传S3失败: {str(e)}"

    def _sample_sync_call_t2v(self, param):
        # 同步调用，直接返回结果
        try:
            print('please wait...')
            rsp = VideoSynthesis.call(api_key=param['api_key'],
                                      model=param['model'],
                                      prompt=param['prompt'],
                                      seed=param['seed'],
                                      size=param['size'])
            print(rsp)
            if rsp.status_code == HTTPStatus.OK:
                if rsp.output.video_url == "":
                    return False, rsp.output.code, 'failed'
                print("video_url:", rsp.output.video_url)
                return True, rsp.output.video_url, 'success'
            else:
                print('Failed, status_code: %s, code: %s, message: %s' %
                      (rsp.status_code, rsp.code, rsp.message))
                return False, ('Failed, status_code: %s, code: %s, message: %s' %
                               (rsp.status_code, rsp.code, rsp.message)), 'failed'
        except InvalidTask as e:
            # 提取 JSON 部分
            error_message = str(e)
            json_start = error_message.find("{")
            json_end = error_message.rfind("}") + 1
            error_data = json.loads(error_message[json_start:json_end])
            error_code = error_data.get("code")
            message = error_data.get("message")
            print(f"错误码: {error_code}")
            print(f"错误信息: {message}")
            print(f"返回值类型: {type(e)}")
            return False, error_code, message

    def _sample_call_i2v(self, param):
        try:
            rsp = VideoSynthesis.call(api_key=param['api_key'],
                                      model=param['model'],
                                      prompt=param['prompt'],
                                      negative_prompt=param.get('negative_prompt', None),
                                      img_url=param['img_url'],
                                      resolution=param['resolution'],
                                      duration=param['duration'],
                                      seed=param['seed']
                                      )
            print("rsp", rsp)
            if rsp.status_code == HTTPStatus.OK:
                if rsp.output.video_url == "":
                    return False, rsp.output.code, 'failed'
                print("video_url:", rsp.output.video_url)
                return True, rsp.output.video_url, 'success'
            else:
                return False, ('Failed, status_code: %s, code: %s, message: %s' %
                               (rsp.status_code, rsp.code, rsp.message)), 'failed'
        except InvalidTask as e:
            # 提取 JSON 部分
            error_message = str(e)
            json_start = error_message.find("{")
            json_end = error_message.rfind("}") + 1
            error_data = json.loads(error_message[json_start:json_end])
            error_code = error_data.get("code")
            message = error_data.get("message")
            print(f"错误码: {error_code}")
            print(f"错误信息: {message}")
            print(f"返回值类型: {type(e)}")
            return False, error_code, message

    def _sample_sync_call_kf2v(self, param):
        print('please wait...')
        try:
            rsp = VideoSynthesis.call(
                api_key=param['api_key'],
                model=param['model'],
                prompt=param['prompt'],
                # negative_prompt=param.get('negative_prompt', None),
                first_frame_url=param['first_frame_url'],
                last_frame_url=param['last_frame_url'],
                resolution=param['resolution'],
                prompt_extend=True)
            print("result:::::", rsp)
            if rsp.status_code == HTTPStatus.OK:
                if rsp.output.video_url == "":
                    return False, rsp.output.code, 'failed'
                print("video_url:", rsp.output.video_url)
                return True, rsp.output.video_url, 'success'
            else:
                return False, ('Failed, status_code: %s, code: %s, message: %s' %
                               (rsp.status_code, rsp.code, rsp.message)), 'failed'
        except InvalidTask as e:
            # 提取 JSON 部分
            error_message = str(e)
            json_start = error_message.find("{")
            json_end = error_message.rfind("}") + 1
            error_data = json.loads(error_message[json_start:json_end])
            error_code = error_data.get("code")
            message = error_data.get("message")
            print(f"错误码: {error_code}")
            print(f"错误信息: {message}")
            print(f"返回值类型: {type(e)}")
            return False, error_code, message
    # -------------------------
    # 外部调用方法：主入口
    # --------------------------
    def process_video_task_wx(self, task_params: Dict[str, Any]):
        """
        处理视频生成任务（外部调用入口）

        参数:
            task_params: 任务参数字典，包含：
                - template : 模板名称
                - taskId: 任务ID
                - moduleName: 模型名称索引
                - prompt: 提示词
                - seed 随机种子 当默认不传或者传0时，
                - modelType: 模型类型（1:文生视频, 2:单图生视频, 3:首尾图生视频, 4:模板, 5:多图特效）
                - duration: 视频时长（默认5）
                - first_image: 首图URL（modelType=2/3/4/5时需要）
                - last_image: 尾图URL（modelType=3/5时需要）
                - storageService: 存储服务类型（如'r2'）
                - bucketName: 存储桶名称
                - bgm:是否为生成的视频添加背景音乐。默认：false，可选值 true 、false
                - storedPrefix: S3存储路径前缀
                - QueueResult: 结果队列名称
                - 其他可选参数：fps, resolution, ratio等

        返回:
            (是否成功, 结果字典) 结果字典包含's3_key'（视频S3路径）或'error'（错误信息）
        """
        try:
            # 1. 参数校验
            required_params = ['prompt', 'modelType', 'storageService', 'bucketName', 'storedPrefix', 'QueueResult']
            for param in required_params:
                if not task_params.get(param):
                    return {"status": 2, "msg": f"缺少参数: {param}", "code": 6003}

            model_type = task_params.get('modelType')
            modelName = task_params.get('moduleName')
            # 2. 构建API请求参数与路径
            api_path = ""
            storage_service = task_params.get('storageService')
            bucket_name = task_params.get('bucketName')
            payload = {'api_key': self.API_KEY, 'model': modelName, "prompt": task_params.get('prompt'), "seed": task_params.get('seed', 0), 'duration': int(task_params.get('duration')), 'negative_prompt': task_params.get('negative_prompt')}
            resolution = task_params.get('resolution').upper()
            # 根据模型类型补充参数

            if model_type == 1:  # 文生视频
                success, res = self.get_resolution(
                    model=task_params.get('moduleName'),
                    quality=resolution,
                    aspect_ratio=task_params.get('ratio')
                )
                if success:
                    payload['size'] = res
                else:
                    return {"status": 2, "msg": res, "code": 6001}
                success, video_result, error_msg = self._sample_sync_call_t2v(payload)
            elif model_type == 2:
                payload['resolution'] = resolution
                first_image_url = task_params.get('first_image')
                if not first_image_url:
                    return {"status": 2, "msg": "modelType=2/4需提供first_image", "code": 6003}
                # first_img = sr_s3.get_url(first_image_url, serviceName=storage_service, bucketName=bucket_name)

                success, img_data, _ = self._download_and_compress_image(
                    first_image_url, task_params['storageService'], task_params['bucketName']
                )
                if not success:
                    return {"status": 2, "msg": img_data, 'code': 6005}
                payload['img_url'] = img_data
                success, video_result, error_msg = self._sample_call_i2v(payload)
            elif model_type == 3:  # 首尾图生视频
                payload['resolution'] = resolution
                first_image_url = task_params.get('first_image')
                last_image_url = task_params.get('last_image')
                if not (first_image_url and last_image_url):
                    return {"status": 2, "msg": "modelType=3/5需提供first_image和last_image", "code": 6003}
                # 下载并处理首图和尾图
                # first_img = sr_s3.get_url(first_image_url, serviceName=storage_service, bucketName=bucket_name)
                #
                # last_img = sr_s3.get_url(last_image_url, serviceName=storage_service, bucketName=bucket_name)
                success, img_data, _ = self._download_and_compress_image(
                    first_image_url, task_params['storageService'], task_params['bucketName']
                )
                if not success:
                    return {"status": 2, "msg": img_data, 'code': 6005}
                success, img_data2, _ = self._download_and_compress_image(
                    last_image_url, task_params['storageService'], task_params['bucketName']
                )
                if not success:
                    return {"status": 2, "msg": img_data, 'code': 6005}
                payload["first_frame_url"] = img_data
                payload["last_frame_url"] = img_data2
                success, video_result, error_msg = self._sample_sync_call_kf2v(payload)
            else:
                return {"status": 2, "msg": f"不支持的模型类型: {model_type}"}
            # 3. 创建视频任务
            # print(payload)
            # exit()
            print(video_result, "video_result")
            if video_result == "DataInspectionFailed" or video_result == "IPInfringementSuspect":
                msg_body = {
                    'status': 7,
                    'width': 0,
                    'height': 0,
                    'local_path': 0,
                    'url': 0,
                    "model_mode": "std",
                    'msg': error_msg,
                    'code': StatusCodes.OUTPUT_VIDEO_INVALID_MODEL
                }
                return msg_body
            elif video_result == 'InvalidParameter.DataInspection':
                msg_body = {
                    'status': 2,
                    'width': 0,
                    'height': 0,
                    'local_path': 0,
                    'url': 0,
                    "model_mode": "std",
                    'msg': error_msg,
                    'code': StatusCodes.IMAGE_PROCESS_ERROR_MODEL
                }
                return msg_body
            if not success:
                return {"status": 2, "msg": video_result, "code": 6001}
            # 5. 下载视频并上传到S3
            success, video_info = self._download_video(video_result, task_params['taskId'])
            if not success:
                return {"status": 2, "msg": video_info.get("error"), "code": 6001}
            print("下载视频地址", video_info["save_path"])
            # 6. 上传视频到S3
            success, s3_key = self._upload_video_to_s3(
                video_info["save_path"],
                task_params['storedPrefix'],
                task_params['storageService'],
                task_params['bucketName']
            )
            if not success:
                return {"status": 2, "msg": s3_key, "code": 6001}

            # 7. 发送结果到队列（假设sr_sqs为队列工具）
            from utils import sr_sqs
            msg_body = {
                'status': 1,
                'width': video_info['width'],
                'height': video_info['height'],
                'local_path': video_info['save_path'],
                'url': s3_key,
                "model_mode": "std",
                'msg': "success",
                'code': ''
            }
            return msg_body

        except Exception as e:
            error_msg = f"任务处理异常: {str(e)}"
            # print(e.status_code)
            return {"status": 2, "msg": error_msg, "code": 6001}


# --------------------------
# 外部调用示例
# --------------------------
if __name__ == "__main__":
    # 初始化视频生成器
    video_generator = wanXiangGenerator_Vid()
    # 示例任务参数
    task_params = {
        'taskId' : '1759026960823ENKPCB4KDBsrngzF',
    'prompt' : 'A cinematic video of a young woman standing under a spotlight, wearing a red checkered shirt over a white T-shirt and a black headscarf. She slowly removes the black scarf, revealing two long thick braids that fall over her chest. Then she unbuttons her checkered shirt, takes it off, and remains in a plain white T-shirt. Next, she undoes her braids, her hair falling free past her shoulders, then tosses her loose hair with a gentle motion and finally smiles warmly. Cinematic lighting, ultra realistic, 4k, emotional transformation from tension to relief.',
    'first_image' : 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1759026934844.png',
    'last_image' : 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1759026943359.jpg',
    'duration' : '5',
    'fps' : '30',
    'resolution' : '720p',
    'ratio' : '',
    'appId' : '1',
    'taskType' : 'pwai_video_itf',
    'QueueResult' : 'AILocalVideoResult',
    'storageService' : 'r2',
    'isLogin' : 1,
    'bucketName' : 'picwand',
    'storedPrefix' : 'video/image_to_video',
    'progress' : 1,
    'queueName' : 'AIHubVideoTasks_Test',
    'modelType' : 3,
    'moduleType' : 27,
    'moduleName' : 'wanx2.1-kf2v-plus',
    'mResolution' : '720p',
    'mode' : 'std',
    'negative_prompt' : '',
    'cfg_scale' : 0.5,
    'p_t_b' : 0,
    'n_t_b' : 0,
    'motion_type' : '',
    'model_style' : '',
    'camera_movement' : '',
    }

    # success = video_compress_image_size_path("../download/1756999547442 (1).avif", 10)
    # print(success)
    # exit()
    # success, err = sr_s3.download_file(
    #     "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1756591704707.heic", "../download/1.heic", serviceName="r2", bucketName="picwand"
    # )
    # print(success,err)
    # exit()
    # task_params = {'taskId': '5387be8f12dddfb862b0', 'prompt': 'Create a square masterpiece video high definition from this image. Seconds 0–2: steam begins to rise from under the wooden benches. At the same time, a thin stream of steam rises continuously from the central stove. Seconds 3–5: The steam from under the benches stops completely. From the perimeter of the ceiling, a heavy tropical rain shower begins. The central stove continues to emit steam as before. Camera does not rotate — it very slowly zooms in toward the central stove. Guests remain relaxed, with only subtle natural movements like shifting posture or tilting their heads.. Atmosphere is serene, immersive, cinematic', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1756948115412.jpg', 'last_image': '', 'duration': '5', 'fps': '30', 'resolution': '480p', 'ratio': '', 'appId': '101', 'taskType': 'pwai_video_itf', 'QueueResult': 'AIHubVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks', 'modelType': 2, 'moduleType': 27, 'moduleName': 'wan2.2-i2v-plus', 'mResolution': '480p', 'mode': 'std', 'negative_prompt': 'Tattoo, slow motion, tattoos, zoom in, close up', 'cfg_scale': 0.5, 'p_t_b': 0, 'n_t_b': 0, 'motion_type': '', 'model_style': '', 'camera_movement': ''}
    # print(task_params)
    result = video_generator.process_video_task_wx(task_params)
    print(result)
